import cv2
from paddleocr import PaddleOCR
import logging
import numpy as np

# 关闭 debug和warning，否则paddleocr会输出很多debug信息。
# Disable DEBUG and INFO logging
logging.disable(logging.DEBUG)  # This line disables all DEBUG logs
# Optionally, you can also disable WARNING logs if needed
logging.disable(logging.WARNING)  # Uncomment this line to disable WARNING logs

# 通过摄像头 cap 拍摄图像，保存到 frame。
def capture(cap):
    ret, frame = cap.read()  # Capture frame from webcam
    return frame


# 对输入的图像 frame 进行字符检测和识别，
# 把检测到的单词保存到 words
def ocr(frame):
    # Paddleocr目前支持的多语言语种可以通过修改lang参数进行切换
    # 例如`ch`, `en`, `fr`, `german`, `korean`, `japan`
    ocr = PaddleOCR(use_angle_cls=True, lang="ch")  # need to run only once to download and load model into memory
    
    result = ocr.ocr(frame, cls=True)
    pieces = []        
    for line in result:
        for word in line:      
            # 有些棋子在 在中文名前面还有数字，
            # 识别出来的单词就会包含3个字符，
            # 这时候只选择最后的2个中文字符。
            # 如果识别出来的单词中字符个数小于2， 表示这个单词不是合格的识别结果。
            if len(word[1][0]) >= 2:
                pieces.append(word[1][0][-2:])
  
    return pieces


def get_ranks(pieces):    
    # 定义用整数表示的棋子等级
    # 1~9表示普通棋子, 10表示炸弹, 0表示军旗, -1表示地雷
    piece_rank_mapping = {
        "司令": 9,
        "军长": 8,
        "师长": 7,
        "旅长": 6,
        "团长": 5,
        "营长": 4,
        "连长": 3,
        "排长": 2,
        "工兵": 1,     
        "地雷": -1,     
        "炸弹": 10,     
        "军旗": 0,     
    }
    piece1 = pieces[0]
    piece2 = pieces[1]
    a = piece_rank_mapping[piece1]
    b = piece_rank_mapping[piece2]
    return a, b


# 判断是否需要亮军旗
def check_flag(a,b,ret):
    _flag={"a":0,"b":0}
    #亮军旗
    if a==9 and (ret==2 or ret==0):
        _flag["a"]=1
    if b==9 and (ret==1 or ret==0):
        _flag["b"]=1
    return _flag


# 对比 A 和 B 的棋子大小
# 返回 result 和 flag
def compare(a,b):
    print(f"a:{a}, b:{b}")
    ret=-1000
    if a<=0 and b<=0:#军旗地雷
        ret=-1
        return ret,check_flag(a,b,ret)        
    
    if a*b==0:#挖军旗
        if a==0:
            ret=4
        else:
            ret=3
        return ret,check_flag(a,b,ret)

    if (a==10 or b==10) or a==b:#炸弹或同大小
        ret=0
        return ret,check_flag(a,b,ret)
    
    if a == -1: # 有地雷
        if b == 1: # 士兵挖地雷
            ret = 2
        else: # 其他棋子碰地雷
            ret = 1
        return ret,check_flag(a,b,ret)

    if b == -1: # 有地雷
        if a == 1: # 士兵挖地雷
            ret = 1
        else: # 其他棋子碰地雷
            ret = 2
        return ret,check_flag(a,b,ret)
        
    #比大小
    if a>b:
        ret=1
    else:
        ret=2

    return ret,check_flag(a,b,ret)

def get_color(result,flag):
        
    red=(0,0,255)
    green=(0,255,0)
    white=(255,255,255)
    black=(0,0,0)
    yellow=(0,255,255)
    blue=(255,0,0)
    
    ca, cb = blue, blue
    
    if result==-1:
        return ca, cb
    if result==0:
        ca=red
        cb=red
    elif result==1:
        ca=green
        cb=red
    elif result==2:
        ca=red
        cb=green
    elif result==3:
        ca=yellow
        cb=black
    elif result==4:
        ca=black
        cb=yellow
    
    if flag["a"]== 1:
        ca=white
    if flag["b"]== 1:
        cb=white
    return ca, cb     

#BGR
# result(int) -1 error, 0 die, 1 a win, 2 b win, 3 a victory, 4 b victory.
# Red die, green win, white flag, yello victory, black lose, blue error.
# 画图显示棋子比较大小的结果
def draw(ca,cb):    
    img = np.zeros((300, 512, 3), np.uint8)
    img_line = img.copy()
    img_line[:,:,:]=[255,255,0]
    ####start code
    cv2.rectangle(img_line,(100,85),(175,175),ca,-1)
    cv2.rectangle(img_line,(100+185,85),(175+185,175),cb,-1)
    font = cv2.FONT_HERSHEY_SIMPLEX
    cv2.putText(img_line, 'A', (100, 170), font,4, (255, 255, 0), 6, lineType=cv2.LINE_AA)
    cv2.putText(img_line, 'B', (280, 170), font,4, (255, 255, 0), 6, lineType=cv2.LINE_AA)
    
    ####end code
    cv2.imshow('img',img_line)
    cv2.waitKey(0)                            
    cv2.destroyAllWindows()  
    
    
